import torch
from torch import nn

from model.conv_bn_relu import ConvBNRelu
from options import HiDDenConfiguration


class Encoder(nn.Module):
    '''
    将数据隐藏进图像中
    '''

    def __init__(self, config: HiDDenConfiguration):
        super(Encoder, self).__init__()
        self.H = config.H
        self.W = config.W
        self.conv_channels = config.encoder_channels
        self.num_blocks = config.encoder_blocks

        '''
        存储编码器的所有基础卷积层的列表
        '''
        layers = [ConvBNRelu(3, self.conv_channels)]

        for _ in range(config.encoder_blocks-1):
            layer = ConvBNRelu(self.conv_channels,self.conv_channels)
            layers.append(layer)

        #基础特征提取网络，将三通道图像转换为encoder_channels,通过多层卷积提取图像特征
        self.conv_layers = nn.Sequential(*layers)

        #加密层，输入通道数为 图像特征通道数+原始图像通道数+秘密信息长度
        self.after_concat_layer = ConvBNRelu(self.conv_channels+3+config.message_length,
                                             self.conv_channels)
        #输出与原图像一致的含加密信息的图像
        self.final_layer = nn.Conv2d(self.conv_channels, 3, kernel_size=1)

    def forward(self, image, message):

        #将原来图片二维[batch_size(批次大小)，message_length(信息长度)]增加两个维度
        expended_message = message.unsqueeze(-1)
        expended_message = expended_message.unsqueeze(-1)

        expended_message = expended_message.expand(-1,-1,self.H,self.W)
        encoded_image = self.conv_layers(image)
        #按维度拼接张量,按第一个维度进行拼接
        concat = torch.cat([expended_message, encoded_image, image], dim=1)

        #将数据嵌入图片中
        im_w = self.after_concat_layer(concat)
        im_w = self.final_layer(im_w)
        return im_w





